import sys

import numpy as np
from mpl_toolkits import mplot3d
from matplotlib import pyplot as plt
import tensorflow as tf
import tensorflow_addons as tfa
from tensorflow.keras import layers, losses, metrics, optimizers
from python_ai.common.xcommon import *
import os

np.random.seed(777)
tf.random.set_seed(777)

sep('X')
X = tf.reshape(tf.range(1, 8 + 1), [4, 2])
tf.print(X, output_stream=sys.stdout, summarize=-1)

sep('Y')
Y = tf.reshape(tf.range(1, 8 + 1), [4, 2])
tf.print(Y, output_stream=sys.stdout, summarize=-1)

sep('squeeze(Y)')
Y_sq = tf.squeeze(Y)
tf.print(Y_sq, output_stream=sys.stdout, summarize=-1)

sep('squeeze(Y) > 2')
Y_sq_lth = tf.squeeze(Y) > 2
tf.print(Y_sq_lth, output_stream=sys.stdout, summarize=-1)

sep('boolean mask')
sep('pos')
X_pos = tf.boolean_mask(X, Y_sq_lth, axis=0)
tf.print(X_pos, output_stream=sys.stdout, summarize=-1)
sep('neg')
X_neg = tf.boolean_mask(X, ~Y_sq_lth, axis=0)
tf.print(X_neg, output_stream=sys.stdout, summarize=-1)

sep('Adjust y_sl_lth')
Y_sq_lth = tf.constant([
    [0, 0],
    [0, 1],
    [1, 1],
    [1, 1],
])
tf.print(Y_sq_lth, output_stream=sys.stdout, summarize=-1)
sep('pos')
X_pos = tf.boolean_mask(X, Y_sq_lth, axis=0)
tf.print(X_pos, output_stream=sys.stdout, summarize=-1)
sep('neg')
X_neg = tf.boolean_mask(X, ~Y_sq_lth, axis=0)
tf.print(X_neg, output_stream=sys.stdout, summarize=-1)
